/*
 *	Copyright © 2008 University of Houston
 *	All rights reserved.
 */
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <stdlib.h>
#include <stdio.h>
#include <iostream>
#include <sstream>
#include "session.h"
#include "logger.h"
#include "filter.h"

using namespace std;

timedHeader::timedHeader()
{
	// Zero out all of the variables
	memset(this,0,sizeof(*this));

	// Get the timestamp
	gettimeofday(&timestamp,NULL);
}

timedHeader::~timedHeader() {};

sniff_tcp *timedHeader::getTCP()
{
	return &tcpHeader;
}
sniff_ip  *timedHeader::getIP()
{
	return &ipHeader;
}

in_addr timedHeader::getSourceIP()
{
	return ipHeader.ip_src;
}
in_addr timedHeader::getDestIP()
{
	return ipHeader.ip_dst;
}
int 	timedHeader::getSourcePort()
{
	return ntohs(tcpHeader.th_sport);
}
int 	timedHeader::getDestPort()
{
	return ntohs(tcpHeader.th_dport);
}
int		timedHeader::getPayloadSize()
{
	return payloadSize;
}
timeval	timedHeader::getTimestamp()
{
	return timestamp;
}
unsigned long	timedHeader::getTime()
{
	return timestamp.tv_sec;
}
bool timedHeader::isFinAck()
{
	if (isFin() && isAck())
		return true;
	return false;
}
bool timedHeader::isAck()
{
	if (getTCP()->th_flags & TH_ACK)
		return true;
	return false;
}
bool timedHeader::isFin()
{
	if (getTCP()->th_flags & TH_FIN)
		return true;
	return false;
}
bool timedHeader::isPush()
{
	if (getTCP()->th_flags & TH_PUSH)
		return true;
	return false;
}
bool timedHeader::seqMatches(timedHeader other)
{
	if (getSequence() == other.getSequence())
		if (getAck() == other.getAck())
			return true;
	return false;
}
bool timedHeader::formatPacket(const u_char* packet)
{
	/* declare pointers to packet headers */
	const struct sniff_ethernet *ethernet;  /* The ethernet header [1] */
	const struct sniff_ip *ip;              /* The IP header */
	const struct sniff_tcp *tcp;            /* The TCP header */

	int size_ip;
	int size_tcp;
	int size_payload;

	/* define ethernet header */
	ethernet = (struct sniff_ethernet*)(packet);

	/* define/compute ip header offset */
	ip = (struct sniff_ip*)(packet + SIZE_ETHERNET);
	size_ip = IP_HL(ip)*4;
	if (size_ip < 20)
		return false;

	if (ip->ip_p != IPPROTO_TCP)
		return false;

	/* define/compute tcp header offset */
	tcp = (struct sniff_tcp*)(packet + SIZE_ETHERNET + size_ip);
	size_tcp = TH_OFF(tcp)*4;
	if (size_tcp < 20)
		return false;

	/* compute tcp payload (segment) size */
	size_payload = ntohs(ip->ip_len) - (size_ip + size_tcp);

	// Populate internal fields.
	memcpy(&tcpHeader, tcp, sizeof(tcpHeader));
	memcpy(&ipHeader, ip, sizeof(ipHeader));
	memcpy(&ethHeader, ethernet, sizeof(ethHeader));
	payloadSize = size_payload;

	return true;
}

void timedHeader::setSrcIP(const char* srcIPBuffer)
{
	inet_aton(srcIPBuffer,&ipHeader.ip_src);
}
void timedHeader::setDstIP(const char* dstIPBuffer)
{
	inet_aton(dstIPBuffer,&ipHeader.ip_dst);
}
void timedHeader::setSrcPort(const u_int srcPort)
{
	tcpHeader.th_sport = htons(srcPort);
}
void timedHeader::setDstPort(const u_int dstPort)
{
	tcpHeader.th_dport = htons(dstPort);
}
void timedHeader::setPayloadSize(const int payload_size)
{
	payloadSize = payload_size;
}
void timedHeader::setFlags(const u_char flags)
{
	getTCP()->th_flags = flags;
}

timedHeader session::operator[](const unsigned int offset) const
{
	if (offset > headerList.size())
	{
		timedHeader empty;
		return empty;
	}

	return headerList[offset];
}

timedHeader& timedHeader::operator=(const u_char* param)
{
	formatPacket(param);
	return *this;
}

timedHeader& timedHeader::operator=(const timedHeader param)
{
	memcpy(this, &param, sizeof(timedHeader));
	return *this;
}
tcp_seq	timedHeader::getSequence()
{
	return getTCP()->th_seq;
}
tcp_seq timedHeader::getAck()
{
	return getTCP()->th_ack;
}

session::session()
{
	// Do nothing
	initialize();
}
session::session(timedHeader packet)
{
	// Add the first packet
	addPacket(packet);

	initialize();
}
session::~session()
{
}
void session::initialize()
{
	// Initialize some variables regarding connection closing
	finack1 = finack2 = final_ack = false;
	finack1_index = 0;
	finack2_index = 0;
}
void session::setAnalyzer(sessionAnalyzer* a)
{
	pAnalyzer = a;
}
int session::getNumPackets() const
{
	return headerList.size();
}
vector<timedHeader>* session::getHeaderList()
{
	return &(headerList);
}

void session::addPacket(timedHeader th)
{
	// Load the packet
	loadPacket(th);

	// Tell the filter that we've gotten a new packet
	if (pAnalyzer)
		pAnalyzer->update();
}

void session::loadPacket(timedHeader th)
{
	headerList.push_back(th);

	if (!th.isAck())
		return;

	if (th.isFin())
	{
		// If we have received the first FIN/ACK packet...
		if (!finack1)
		{
			// Keep track of what the index is for this packet
			finack1_index = headerList.size()-1;
			finack1 = true;
		}
		// We have received the first FIN/ACK packet.  Check to see if this
		// is the FIN/ACK response.  Note that we have to make sure that the
		// hosts have been reversed since the first FIN/ACK, to ensure that this
		// is not just a re-send of the first packet.
		else if (!finack2)
		{
			timedHeader th_finack1 = headerList[finack1_index];

			// Make sure the packets go together
			if (th_finack1.getAck() == th.getSequence())
			{
				finack2_index = headerList.size()-1;
				finack2=true;
			}
		}
	}

	// The FIN/ACK handshake is almost complete.  Just need to check on the
	// final ACK packet.
	else if (finack2 && !final_ack)
	{
		timedHeader th_finack2 = headerList[finack2_index];

		if (th_finack2.getAck() == th.getSequence())
			final_ack = true;
	}
}

timedHeader	session::getPacket(int offset) const
{
	return (*this)[offset];
}

in_addr session::getIPA() const
{
	return getPacket(0).getSourceIP();
}
in_addr session::getIPB() const
{
	return getPacket(0).getDestIP();
}
u_short	session::getPortA() const
{
	return getPacket(0).getSourcePort();
}
u_short	session::getPortB() const
{
	return getPacket(0).getDestPort();
}
bool session::belongs(timedHeader th) const
{
	// Get all of the information on the IP addresses
	in_addr srcIP 	= th.getSourceIP();
	in_addr dstIP 	= th.getDestIP();
	in_addr ipA		= getIPA();
	in_addr ipB		= getIPB();

	// Get all of the information on the ports
	u_short srcPrt	= th.getSourcePort();
	u_short dstPrt	= th.getDestPort();
	u_short portA	= getPortA();
	u_short portB	= getPortB();

	// Check to see if the packet is travelling from ipA:portA -> ipB:portB
	if (srcIP == ipA && dstIP == ipB)
		if (srcPrt == portA && dstPrt == portB)
			return true;
	// Check to see if the packet is travelling from ipB:portB -> ipA:portA
	if (srcIP == ipB && dstIP == ipA)
		if (srcPrt == portB && dstPrt == portA)
			return true;

	// The packet does not belong
	return false;
}

bool session::isConnectionClosed() const
{
	return final_ack;
}

timedHeader session::getLatestPacket() const
{
	return getPacket(getNumPackets()-1);
}
int session::getSessionID() const
{
	return sessionID;
}
void session::setSessionID(int i)
{
	sessionID = i;
}

